#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Qwen3-0.6B 简单测试示例
"""

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import os

def simple_test():
    """
    简单测试Qwen3-0.6B模型
    """
    model_path = "/home/zkh/lzt/damoxing/boshuyuce/xindaoyuce/Qwen3-0.6B"
    
    print("正在加载Qwen3-0.6B模型...")
    
    # 加载tokenizer和模型
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        device_map="auto" if torch.cuda.is_available() else None,
        trust_remote_code=True
    )
    
    print(f"模型加载完成！")
    print(f"设备: {next(model.parameters()).device}")
    print(f"数据类型: {next(model.parameters()).dtype}")
    
    # 测试文本生成
    test_prompts = [
        "你好，我是",
        "人工智能是",
        "今天天气",
        "Python编程语言的特点是"
    ]
    
    print("\n开始测试文本生成...")
    print("=" * 50)
    
    for i, prompt in enumerate(test_prompts, 1):
        print(f"\n测试 {i}: {prompt}")
        print("-" * 30)
        
        # 编码输入
        inputs = tokenizer.encode(prompt, return_tensors="pt")
        device = next(model.parameters()).device
        inputs = inputs.to(device)
        
        # 生成文本
        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_length=inputs.shape[1] + 50,  # 生成50个新token
                temperature=0.7,
                top_p=0.9,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                repetition_penalty=1.1
            )
        
        # 解码输出
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        new_text = generated_text[len(prompt):].strip()
        
        print(f"输入: {prompt}")
        print(f"输出: {new_text}")

if __name__ == "__main__":
    try:
        simple_test()
        print("\n✅ 测试完成！")
    except Exception as e:
        print(f"❌ 测试失败: {e}")
        import traceback
        traceback.print_exc() 